#include <zephyr.h>

#include <stdio.h>
#include <string.h>

#include "mbedtls/rsa.h"
#include "mbedtls/cipher.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/platform.h"

#define assert_exit(cond, ret) \
    do { if (!(cond)) { \
        printf("  !. assert: failed [line: %d, error: -0x%04X]\n", __LINE__, -ret); \
        goto cleanup; \
    } } while (0)

typedef struct cipher
{
  uint8_t cipher_pre[2048/8];
  uint8_t cipher_back[2048/8];
}cipher_str;


typedef struct msg_sign
{
  uint8_t sign_pre[2048/8];
  uint8_t msg_back[2048/8];
}msg_sign_str;


cipher_str cipher,cipher_bak;
msg_sign_str msg_sig,msg_sig_bak;



 uint8_t key[32] =
{
    0x06, 0xa9, 0x21, 0x40, 0x36, 0xb8, 0xa1, 0x5b,
    0x51, 0x2e, 0x03, 0xd5, 0x34, 0x12, 0x00, 0x06
};

uint8_t iv[32] =
{
    0x3d, 0xaf, 0xba, 0x42, 0x9d, 0x9e, 0xb4, 0x30,
    0xb4, 0x22, 0xda, 0x80, 0x2c, 0x9f, 0xac, 0x41
};

uint8_t  buf0[256];
uint8_t  buf1[256];


static void dump_buf(char *info, uint8_t *buf, uint32_t len)
{
    mbedtls_printf("%s", info);
    for (int i = 0; i < len; i++) {
        mbedtls_printf("%s%02X%s", i % 16 == 0 ? "\n     ":" ",
                        buf[i], i == len - 1 ? "\n":"");
    }
}

static int entropy_source(void *data, uint8_t *output, size_t len, size_t *olen)
{
    uint32_t seed;

    seed = sys_rand32_get();
    if (len > sizeof(seed)) {
        len = sizeof(seed);
    }

    memcpy(output, &seed, len);

    *olen = len;
    return 0;
}


void cipher_decrypt(int type,uint8_t *input)
{

    size_t len;
    int olen = 0;


    mbedtls_cipher_context_t ctx;
    const mbedtls_cipher_info_t *info;


    mbedtls_platform_set_printf(printf);

    mbedtls_cipher_init(&ctx);
    info = mbedtls_cipher_info_from_type(type);

    mbedtls_cipher_setup(&ctx, info);

    //解密
    mbedtls_cipher_setkey(&ctx,key, sizeof(key)*16, MBEDTLS_DECRYPT);

    mbedtls_cipher_crypt(&ctx,iv,sizeof(iv),input, strlen(input),buf0, &len );
    olen += len;

    mbedtls_cipher_finish(&ctx, buf0 + len, &len);

    olen += len;

    printf("in in ni 456");
     //printf("\n %s",buf0);
    dump_buf("\n  cipher aes decrypt:", buf0, olen);

    mbedtls_cipher_free(&ctx);


}


void cipher_encrypt(int type,uint8_t *input)
{
    size_t len;
    int olen = 0;


    mbedtls_cipher_context_t ctx;
    const mbedtls_cipher_info_t *info;

    mbedtls_platform_set_printf(printf);

    mbedtls_cipher_init(&ctx);
    info = mbedtls_cipher_info_from_type(type);

    mbedtls_cipher_setup(&ctx, info);


    mbedtls_cipher_setkey(&ctx, key, sizeof(key)*8, MBEDTLS_ENCRYPT);
    mbedtls_cipher_update(&ctx, input, strlen(input), buf1, &len);
    olen += len;

    mbedtls_cipher_finish(&ctx, buf1 + len, &len);
    olen += len;


   mbedtls_printf("\n 3. AES密文写入成功\n");

    dump_buf("\n  cipher aes encrypt:",buf1, olen);

    mbedtls_cipher_free(&ctx);
}



int main(void)
{
      int ret;
    size_t olen = 0;


    mbedtls_platform_set_printf(printf);
    mbedtls_platform_set_snprintf(snprintf);

    mbedtls_rsa_context ctx;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    const char *pers = "simple_rsa";
    const char *msg = "Hello, World!";

    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_rsa_init(&ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);

    mbedtls_entropy_add_source(&entropy, entropy_source, NULL,
                               MBEDTLS_ENTROPY_MAX_GATHER,
                               MBEDTLS_ENTROPY_SOURCE_STRONG);
    ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                    (const uint8_t *) pers, strlen(pers));
    assert_exit(ret == 0, ret);
    mbedtls_printf("\n  . setup rng ... ok\n");



    mbedtls_printf("\n  ! RSA Generating large primes may take minutes! \n");
    ret = mbedtls_rsa_gen_key(&ctx, mbedtls_ctr_drbg_random,
                                        &ctr_drbg, 2048, 65537);
    assert_exit(ret == 0, ret);
    mbedtls_printf("\n  1. RSA generate key ... ok\n");


        //1.将消息哈希并用公钥签名     H=Hash(M)   S=Sig(H,Kda)
    ret = mbedtls_rsa_pkcs1_sign(&ctx, mbedtls_ctr_drbg_random, &ctr_drbg,
                                    MBEDTLS_RSA_PRIVATE, MBEDTLS_MD_SHA256,
                                    sizeof(msg), msg,msg_sig.sign_pre);
   assert_exit(ret == 0, ret);
    dump_buf("  2. 将消息哈希并签名:", msg_sig.sign_pre, sizeof(msg_sig.sign_pre));
    dump_buf("\n  msg :", msg, sizeof(msg));

    strcpy(msg_sig.msg_back,msg);

    printf("\n sign and msg were loaded struct \n");

    uint8_t  msg_hash_sig[sizeof(msg_sig.sign_pre)+sizeof(msg_sig.msg_back)];

    memcpy(msg_hash_sig,&msg_sig,sizeof(msg_sig));
/*
    dump_buf("\n  msg and sign", msg_hash_sig, sizeof(msg_hash_sig));

    memcpy(&msg_sig_bak,msg_hash_sig,sizeof(msg_sig_bak));
    dump_buf("\n  sign :", msg_sig_bak.sign_pre, sizeof(msg_sig_bak.sign_pre));
    dump_buf("\n  msg_bak :", msg_sig_bak.msg_back,sizeof(msg_sig_bak.msg_back));

    printf("\n %s \n",msg_sig_bak.msg_back);
*/
    cipher_encrypt(MBEDTLS_CIPHER_AES_256_CTR,msg_hash_sig);

    uint8_t key_bak[32];
    strcpy(key_bak,key);

    printf("sucess3465");
    cipher_decrypt(MBEDTLS_CIPHER_AES_256_CTR,buf0);

    printf("sucess");

    //4 .RSA加密AES公钥         C1 = E（K，Keb）
    ret = mbedtls_rsa_pkcs1_encrypt(&ctx, mbedtls_ctr_drbg_random,
                                    &ctr_drbg, MBEDTLS_RSA_PUBLIC, sizeof(key), cipher.cipher_back, cipher.cipher_back);
    assert_exit(ret == 0, ret);
    dump_buf("\n  4. RSA 加密AES公钥完成", cipher.cipher_back, sizeof(cipher.cipher_back));

    /**************************************************************************************/

    uint8_t key_d[32] = {0};

    ret = mbedtls_rsa_pkcs1_decrypt(&ctx, mbedtls_ctr_drbg_random, &ctr_drbg,
                                    MBEDTLS_RSA_PRIVATE, &olen,cipher.cipher_back , key_d, sizeof(key_d));
    assert_exit(ret == 0, ret);

    key_d[olen] = 0;
    mbedtls_printf("\n  3. RSA decryption ... ok\n     %s\n", key_d);

    if( strcmp(key_d,key) == 1)
    {
        printf("\n right \n");
    }
    else
    {
        printf("\n wrong \n");
    }

















cleanup:
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);
    mbedtls_rsa_free(&ctx);

    return ret;

   

}
